//
//  WaterfallFlowLayout.m
//  XMJJ_zml
//
//  Created by 杨何华 on 2025/5/8.
//

#import "WaterfallFlowLayout.h"
@interface WaterfallFlowLayout ()

// 存放每个 item 的布局属性
@property (nonatomic, strong) NSMutableArray<UICollectionViewLayoutAttributes *> *attributesArray;

// 存放每列的当前高度
@property (nonatomic, strong) NSMutableArray<NSNumber *> *columnHeights;

@end

@implementation WaterfallFlowLayout

- (void)prepareLayout {
    [super prepareLayout];

    // 初始化属性数组
    self.attributesArray = [NSMutableArray array];

    // 初始化列高度数组
    self.columnHeights = [NSMutableArray arrayWithCapacity:self.columnCount];
    for (int i = 0; i < self.columnCount; i++) {
        [self.columnHeights addObject:@(self.sectionInset.top)];
    }

    // 获取 item 数量
    NSInteger itemCount = [self.collectionView numberOfItemsInSection:0];

    // 计算每个 item 的布局属性
    for (NSInteger i = 0; i < itemCount; i++) {
        NSIndexPath *indexPath = [NSIndexPath indexPathForItem:i inSection:0];
        UICollectionViewLayoutAttributes *attributes = [UICollectionViewLayoutAttributes layoutAttributesForCellWithIndexPath:indexPath];

        // 计算 item 的宽度
        CGFloat collectionViewWidth = self.collectionView.frame.size.width;
        CGFloat availableWidth = collectionViewWidth - self.sectionInset.left - self.sectionInset.right - (self.columnCount - 1) * self.columnSpacing;
        CGFloat itemWidth = availableWidth / self.columnCount;

        // 获取 item 的高度
        CGFloat itemHeight = [self.delegate collectionView:self.collectionView layout:self heightForItemAtIndexPath:indexPath];

        // 找到当前最短的列
        NSInteger shortestColumnIndex = 0;
        CGFloat shortestColumnHeight = [self.columnHeights[0] floatValue];
        for (NSInteger j = 1; j < self.columnCount; j++) {
            CGFloat currentColumnHeight = [self.columnHeights[j] floatValue];
            if (currentColumnHeight < shortestColumnHeight) {
                shortestColumnIndex = j;
                shortestColumnHeight = currentColumnHeight;
            }
        }

        // 计算 item 的位置
        CGFloat x = self.sectionInset.left + (itemWidth + self.columnSpacing) * shortestColumnIndex;
        CGFloat y = shortestColumnHeight + self.rowSpacing;

        // 设置 item 的布局属性
        attributes.frame = CGRectMake(x, y, itemWidth, itemHeight);
        [self.attributesArray addObject:attributes];

        // 更新列高度
        self.columnHeights[shortestColumnIndex] = @(y + itemHeight);
    }
}

- (NSArray<UICollectionViewLayoutAttributes *> *)layoutAttributesForElementsInRect:(CGRect)rect {
    NSMutableArray<UICollectionViewLayoutAttributes *> *visibleAttributes = [NSMutableArray array];
    for (UICollectionViewLayoutAttributes *attributes in self.attributesArray) {
        if (CGRectIntersectsRect(attributes.frame, rect)) {
            [visibleAttributes addObject:attributes];
        }
    }
    return visibleAttributes;
}

- (UICollectionViewLayoutAttributes *)layoutAttributesForItemAtIndexPath:(NSIndexPath *)indexPath {
    return self.attributesArray[indexPath.item];
}

- (CGSize)collectionViewContentSize {
    CGFloat maxColumnHeight = 0;
    for (NSNumber *height in self.columnHeights) {
        if (height.floatValue > maxColumnHeight) {
            maxColumnHeight = height.floatValue;
        }
    }
    return CGSizeMake(self.collectionView.frame.size.width, maxColumnHeight + self.sectionInset.bottom);
}

@end
